import os
import math
import numpy as np

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
from matplotlib import rcParams
from scipy.signal import savgol_filter

line_styles = ['-', '--', '-.', ':', 'None', ' ', '', 'solid', 'dashed', 'dashdot', 'dotted']
mark_styles = ["", "o", "v", "*", "^", "+", "D", "x"]
line_colors = ["k", "b", "r", "g", "y", "c", "grey"]

rcParams['font.family'] = 'sans-serif'
rcParams['font.sans-serif'] = ['SimHei']
rcParams['font.sans-serif'] = ['MicroSoft YaHei']
rcParams['font.size'] = 14
rcParams['ytick.labelsize'] = 14
rcParams['xtick.labelsize'] = 14


def weighted_least_square_smooth(result, input_size, stride_size):
    if len(result.shape) > 1:
        result = result.reshape(-1)
    length = result.shape[0]
    rul_curve = np.zeros_like(result)
    for i in range(length):
        t = input_size + np.arange(0, i + 1, 1) * stride_size
        weight = t / np.sum(t)
        R = np.sum(weight * result[:i + 1])
        rul_curve[i] = R
    return rul_curve


def kalman_smooth(result, q=0.0001, r=0.012):
    if len(result.shape) > 1:
        result = result.reshape(-1)
    n = len(result)
    if n <= 1:
        return result
    kalman_result = np.zeros_like(result)
    x0 = result[0]  # 令第一个估计值，为当前值
    p0 = 1.0
    kalman_result[0] = x0
    for i in range(1, n):  # kalman 滤波实时计算，只要知道当前值z就能计算出估计值(后验值)x0
        # 先验值
        x1_minus = x0  # X(k|k-1) = AX(k-1|k-1) + BU(k) + W(k), A=1,BU(k) = 0
        p1_minus = p0 + q  # P(k|k-1) = AP(k-1|k-1)A' + Q(k), A=1
        # 更新K和后验值
        k1 = p1_minus / (p1_minus + r)  # Kg(k)=P(k|k-1)H'/[HP(k|k-1)H' + R], H=1
        x0 = x1_minus + k1 * (result[i] - x1_minus)  # X(k|k) = X(k|k-1) + Kg(k)[Z(k) - HX(k|k-1)], H=1
        p0 = (1 - k1) * p1_minus  # P(k|k) = (1 - Kg(k)H)P(k|k-1), H=1
        kalman_result[i] = x0
    return kalman_result


def sg_smooth(result, window_length=31):
    assert window_length % 2 == 1
    if len(result.shape) > 1:
        result = result.reshape(-1)

    return savgol_filter(result, window_length, 1)


def moving_average_smooth(result, window=3):
    if len(result.shape) > 1:
        result = result.reshape(-1)
    ma_result = np.zeros_like(result)
    assert window % 2 == 1
    t = window // 2
    for i in range(len(ma_result)):
        if i < t:
            ma_result[i] = result[i]
        else:
            ma_result[i] = np.mean(result[i - t: i + t])

    return ma_result


def plot_iteration(num_epoch, loss_dict):
    n = len(loss_dict)
    ep = np.arange(num_epoch)
    plt.plot()
    plt.title("迭代次数")
    c = 0
    for name, curve in loss_dict.items():
        color = line_colors[c % len(line_colors)]
        style = line_styles[c % len(line_styles)]
        plt.plot(ep, curve, color=color, linestyle="-")
        c += 1
    plt.show()


# def compare_predict_method(result, label, input_size, stride_size):
#     length = result.shape[0]
#     x = np.arange(0, length, 1)
#     t = input_size + x * stride_size
#     r1 = []
#     r2 = []
#     for l in range(1, length):
#         r1.append(weighted_least_square_predict(result[:l], input_size, stride_size)[-1])
#         r2.append(normal_predict(result[:l], input_size, stride_size)[-1])

#     plt.figure()
#     plt.title("预测方向")
#     plt.plot(x, label, color="red")
#     plt.scatter(x[1:], r1, color="blue")
#     plt.scatter(x[1:], r2, color="green")
#     plt.show()


def plot_predictions(results, labels, names):
    assert len(results) == len(labels)
    assert len(results) == len(names)
    n = len(results)
    zs = [i for i in range(n)]
    ax = plt.figure().add_subplot(projection="3d")
    plt.title("hello world")
    for i in range(n):
        z = zs[i]
        x = np.arange(results[i].shape[0])
        ax.plot(x, results[i], z)
        ax.plot(x, labels[i], z)
    # 图像的其他参数设置
    ax.view_init(20, 45)  # 图形展示角度
    # ax.grid(False)#去掉网格
    ax.w_xaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))  # 背景设置为白色
    ax.w_yaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
    ax.w_zaxis.set_pane_color((1.0, 1.0, 1.0, 1.0))
    # 坐标及坐标轴相关设置
    ax.set_xlabel('x')  # x轴坐标名称及字体样式
    ax.set_ylabel('y')  # x轴坐标名称及字体样式
    ax.set_zlabel('z')  # z轴坐标名称及字体样式
    # ax.set_xlim(0,20)#x轴范围
    # ax.set_zlim(0,40)
    ax.set_yticks([1, 2, 3, 4, 5])  # y轴刻度字体大小
    plt.show()


def calculate_score(results, labels):
    eps = 1e-6
    n = len(results)
    score = 0
    ers = []
    # print(f"labels[n-1] = {labels[-1]}, results[n-1] = {results[n-1]}")
    # er = 100  * (labels[0][-1] - results[0][-1]) / (labels[0][-1] + eps)
    # c = 5 if er <= 0 else 20
    # v = -1 if er <= 0 else 1
    # ai = math.exp(v * math.log(0.5) * er / c)

    for i in range(n):
        m = len(results[i])
        # p = results[i] if results[i] > 0 else 0
        p = results[i]
        l = labels[i]
        ers = 100 * (l - p) / (l + eps)
        for er in ers:
            if er < -50:
                er = -50
            ai = np.exp(np.sign(er) * np.log(0.5) * er / (np.sign(er) * 7.5 + 12.5))
            score += ai
        score /= m

    # for i in range(n):
    #     p = results[i][-1] if results[i][-1] > 0 else 0
    #     l = labels[i][-1]
    #     er = 100 * (l - p) / (l + eps)
    #     ers.append(er)
    #     c = 5 if er <= 0 else 20
    #     v = -1 if er <= 0 else 1
    #     ai = math.exp(v * math.log(0.5) * er / c)
    #     score += ai
    return ers, score / n


def plot_results(result, label, stride_size, names, save_path, split_ratio=0, full_life=1):
    n = len(result)

    for m in range(n):
        result[m] = np.maximum(result[m], 0)
        seq_len = result[m].shape[0]
        bi = int(seq_len * split_ratio)
        # curve = kalman_smooth(result[m])
        error = result[m] - label[m] * full_life
        x_axis = np.zeros(seq_len)

        np.save(os.path.join(save_path, names[m] + "_result.npy"), result[m])
        np.save(os.path.join(save_path, names[m] + "_label.npy"), label[m])
        # np.save(os.path.join(save_path, names[m] + "_smooth.npy"), curve)

        plt.figure(figsize=(8, 5))

        t = np.arange(seq_len - bi) * stride_size
        q = seq_len - 800
        t1 = np.arange(seq_len - 800, seq_len) * stride_size
        plt.plot(t, result[m][bi:], label="预测RUL", c="b", linestyle='-', linewidth=2)
        plt.plot(t, result[m][bi:], color="b", linestyle='-', alpha=0.7)

        plt.plot(t, label[m][bi:], label="轴承真实RUL", c="r", linestyle='--', linewidth=2)
        plt.plot(t, label[m][bi:], color="r", linestyle='--', alpha=0.7)

        plt.bar(t, error, label="误差", color='#77007F')
        plt.plot(t, x_axis, color='#77007F')
        plt.xlim(0)

        # plt.plot(t1, result[m][-800:], label="预测RUL", c="b", linestyle='-', linewidth=2)
        # # plt.plot(t1, label[m][-800:], color="b", linestyle='-', alpha=0.7)
        # plt.plot(t1, label[m][-800:] * full_life, label="轴承真实RUL", c="r", linestyle='--', linewidth=2)
        # # 添加误差填充区域
        # plt.bar(t1, error[-800:], label="误差", color='#77007F')
        # plt.plot(t1, x_axis[-800:], color='#77007F')
        # plt.xlim(seq_len - 800, seq_len)

        # plt.plot(t1, curve[q:] * full_life, label="卡尔曼滤波曲线", c="b")

        plt.xlabel("时间步")
        plt.ylabel("剩余寿命")
        plt.title(names[m])
        ax = plt.gca()
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        # plt.ylim(-0.1,1)
        # yticks = np.arange(0, 1.1, 0.5)
        # plt.yticks(yticks)
        # 添加 RMSE 标签
        # plt.text(seq_len-750, -0.1, 'RMSE=0.0143', fontsize=12, color='black')
        # plt.text(100, -0.2, 'RMSE=0.0416', fontsize=12, color='black')

        plt.tight_layout()
        plt.legend()
        plt.savefig(os.path.join(save_path, names[m] + ".png"))
        plt.show()
        plt.close()
